import { config } from "dotenv";

// 强制加载.env文件，覆盖系统环境变量
config({ override: true });

import MCPClient from "./MCPClient";
import Agent from "./Agent";
import path from "path";
import EmbeddingRetriever from "./EmbeddingRetriever";
import fs from "fs";
import { logTitle } from "./utils";

const URL = "https://news.ycombinator.com/";
const outPath = path.join(process.cwd(), "output");
const TASK = `
我需要你完成一个数据导出任务：查询SQLite数据库中的所有数据并导出为Excel文件。请严格按照以下步骤执行：

**第一步：探索SQLite数据库结构**
- 使用SQLite的list_tables工具查看数据库中的所有表
- 使用describe_table工具获取每个表的结构信息
- 分析表的字段类型和数据结构

**第二步：查询所有表的数据**
- 对每个表使用read_query工具查询所有数据
- 统计每个表的记录数量
- 获取完整的数据内容用于导出

**第三步：数据格式化处理**
- 将查询结果转换为Excel兼容的格式
- 处理特殊字符和数据类型转换
- 为每个表准备CSV格式的数据（Excel可以直接打开CSV）

**第四步：生成Excel文件**
- 使用write_file工具为每个表创建单独的CSV文件
- 文件命名格式：{表名}_data.csv
- 确保CSV文件包含表头和所有数据行
- 保存到output目录

**第五步：创建数据导出报告**
- 生成数据导出汇总报告，包括：
  * 数据库表统计信息
  * 每个表的记录数量
  * 导出文件列表
  * 数据完整性验证结果
- 使用write_file工具将报告保存为"SQLite数据导出报告.md"

**执行要求：**
1. 确保导出数据的完整性和准确性
2. 处理好特殊字符和编码问题
3. 每完成一个步骤，请明确说明完成状态和数据统计
4. 如果某个工具调用失败，请说明原因并尝试替代方案
5. 最终提供完整的导出总结报告
6. CSV文件要包含UTF-8 BOM以确保中文字符正确显示
7. 确保所有工具调用都成功，并记录每个工具调用的耗时和结果

请立即开始执行第一步：探索SQLite数据库中的表结构。

任务完成后就直接输出结果，使用markdown格式输出。

`;
const fetchMCP = new MCPClient("mcp-server-fetch", "uvx", ["mcp-server-fetch"]);
const fileMCP = new MCPClient("mcp-server-file", "npx", [
  "-y",
  "@modelcontextprotocol/server-filesystem",
  process.cwd(),
]);
const elasticsearchMCP = new MCPClient(
  "elasticsearch-mcp-server",
  "uvx",
  ["elasticsearch-mcp-server"],
  {
    ELASTICSEARCH_HOSTS: "https://localhost:9200",
    ELASTICSEARCH_USERNAME: "elastic",
    ELASTICSEARCH_PASSWORD: "test123",
  }
);
const sqliteMCP = new MCPClient("mcp-server-sqlite", "uvx", [
  "mcp-server-sqlite",
  "--db-path",
  "./test.db",
]);

async function main() {
  // RAG
  const context = await retrieveContext();

  // Agent
  const agent = new Agent(
    "Qwen/Qwen3-8B",
    [fetchMCP, fileMCP, elasticsearchMCP, sqliteMCP],
    "",
    context
  );

  // 配置提示词纠正功能
  agent.setPromptCorrectionEnabled(true); // 启用提示词纠正
  agent.setMaxFailures(2); // 设置失败阈值为2次

  logTitle("PROMPT CORRECTION ENABLED");
  process.stdout.write("✅ 提示词自动纠正功能已启用");
  process.stdout.write("📊 失败阈值: 2次");
  process.stdout.write("🔧 当工具连续失败时将自动分析并优化提示词\n");

  await agent.init();

  try {
    await agent.invoke(TASK);
  } catch (error) {
    console.error("任务执行失败:", error);

    // 显示失败统计
    const failureStats = agent.getFailureStats();
    if (failureStats.size > 0) {
      logTitle("FAILURE STATISTICS");
      for (const [toolName, count] of failureStats) {
        process.stdout.write(`${toolName}: ${count} 次失败`);
      }
    }
  } finally {
    await agent.close();
  }
}

main();

async function retrieveContext() {
  // RAG
  const embeddingRetriever = new EmbeddingRetriever("BAAI/bge-m3");
  const knowledgeDir = path.join(process.cwd(), "knowledge");
  const files = fs.readdirSync(knowledgeDir);
  for await (const file of files) {
    const content = fs.readFileSync(path.join(knowledgeDir, file), "utf-8");
    await embeddingRetriever.embedDocument(content);
  }
  const context = (await embeddingRetriever.retrieve(TASK, 3)).join("\n");
  logTitle("CONTEXT");
  process.stdout.write(context);
  return context;
}
